import simnibs
import os
import subprocess
import platform
from simnibs.utils import transformations
import multiprocessing
import argparse
from datetime import datetime

# Post-processing functions
def post_process_silent(opt_name, subject_folder):
    try:
        transformations.middle_gm_interpolation(
            opt_name + ".msh", subject_folder, opt_name + '_subject_overlays',
            out_fsaverage=opt_name + '_fsavg_overlays', depth=0.5,
            open_in_gmsh=False, f_geo=opt_name + "_el_currents.geo"
        )
    except:
        print("middle_gm_interpolation failed")
    try:
        transformations.interpolate_to_volume(
            opt_name + ".msh", subject_folder, opt_name, keep_tissues=[1, 2, 3, 6]
        )
    except:
        print("interpolate_to_volume failed")

    try:
        transformations.warp_volume(
            opt_name + ".msh", subject_folder, opt_name, keep_tissues=[1, 2, 3, 6]
        )
    except:
        print("warp_volume failed")

def post_process(opt_name, subject_folder):
    post_process_silent(opt_name, subject_folder)
    module = "itk-snap" if platform.system() == 'Windows' else "itksnap"
    if os.path.exists(f"{opt_name}_magnE.nii.gz"):
        subprocess.run([module, '-g', f'{opt_name}_magnE.nii.gz'])
    if os.path.exists(f"{opt_name}_MNI_magnE.nii.gz"):
        subprocess.run([module, '-g', f'{opt_name}_MNI_magnE.nii.gz'])


# Simulation function
def run_simulation(subject_folder, target_coord,output_path=None, ldfd_hdf=None,ldfd_hdf_volume=False,target_coord_in_mni=False,
                   target_intensity=None, avoid_coord=None, avoid_coord_in_mni=False, silent=False,
                   max_total_current=2e-3,max_individual_current=1e-3,max_active_electrodes=8):
    if ldfd_hdf==None:
        ldfd_file=os.path.basename(subject_folder)+"_leadfield_EEG10-10_UI_Jurak_2007.hdf5"
        ldfd_hdf=os.path.join(subject_folder,"leadfield",ldfd_file)
        print("Using default path:\n",ldfd_hdf)
    if not os.path.exists(ldfd_hdf):
        print("leadfield_hdf not exists, check if exits or wrong path:\n",ldfd_hdf)
    if output_path==None:
        optimize_dir=os.path.join(subject_folder,"optimize")
        if not os.path.exists(optimize_dir):
            os.mkdir(optimize_dir)
        output_path=os.path.join(optimize_dir,"optimize_results"+datetime.now().strftime('%Y%m%d%H%M%S'))
        print("Using default output_path:\n",output_path)
    if target_intensity==None:
        target_intensity=0.2
        print("Using default intensity for target:\n",target_intensity)
    
    opt = simnibs.opt_struct.TDCSoptimize()
    opt.leadfield_hdf = ldfd_hdf
    #opt.name = output_path
    os.makedirs(output_path,exist_ok=True)
    opt.name = os.path.join(output_path,"optimize")
    opt.max_total_current = max_total_current
    opt.max_individual_current = max_individual_current
    opt.max_active_electrodes = max_active_electrodes
    
    # Define the target
    target = opt.add_target()
    target.positions = target_coord
    target.positions = simnibs.mni2subject_coords(target_coord, subject_folder) if target_coord_in_mni else target_coord
    target.intensity = target_intensity
    
    if ldfd_hdf_volume:
        target.directions = None  # Avoid normal direction error for volumetric data
    
    if avoid_coord:
        avoid = opt.add_avoid()
        avoid.positions = simnibs.mni2subject_coords(avoid_coord, subject_folder) if avoid_coord_in_mni else avoid_coord
    
    # Run optimization
    simnibs.run_simnibs(opt)
    
    # Run post-process functions in parallel
   
    process_func = post_process_silent if silent else post_process
    p = multiprocessing.Process(target=process_func, args=(opt.name, subject_folder))
    p.start()

# Main function with argument parsing
def main():
    parser = argparse.ArgumentParser(description="Run tDCS optimization using SimNIBS")
    parser.add_argument('--ldfd_hdf', help="Path to leadfield HDF file")
    parser.add_argument('--output_path', help="Output path for simulation results")
    parser.add_argument('--target_coord', required=True, type=float, nargs=3, help="Target coordinates,for example, -37, -21, 58 in MNI space")
    parser.add_argument('--subject_folder', required=True, help="Subject folder path")
    parser.add_argument('--target_coord_in_mni', action='store_true', help="Set if target coordinates are in MNI space")
    parser.add_argument('--target_intensity', type=float, default=0.2, help="Target intensity (default: 0.2)")
    parser.add_argument('--avoid_coord', type=float, nargs=3, help="Avoidance coordinates")
    parser.add_argument('--avoid_coord_in_mni', action='store_true', help="Set if avoidance coordinates are in MNI space")
    parser.add_argument('--silent', action='store_true', help="Run in silent mode (no visualization)")
    parser.add_argument('--ldfd_hdf_volume', action='store_true', help="Indicates if the leadfield HDF is volumetric")
    parser.add_argument('--max_total_current', help="Indicates max_total_current")
    parser.add_argument('--max_individual_current', help="Indicates max_individual_current")
    parser.add_argument('--max_active_electrodes',  help="Indicates max_active_electrodes'")
    
    args = parser.parse_args()
    
    run_simulation(
        ldfd_hdf=args.ldfd_hdf,
        output_path=args.output_path,
        target_coord=args.target_coord,
        subject_folder=args.subject_folder,
        target_coord_in_mni=args.target_coord_in_mni,
        target_intensity=args.target_intensity,
        avoid_coord=args.avoid_coord,
        avoid_coord_in_mni=args.avoid_coord_in_mni,
        silent=args.silent,
        ldfd_hdf_volume=args.ldfd_hdf_volume
    )


if __name__ == "__main__":
    main()
